# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import math

import coin.exchange.base.kr_rest.private_client_base as privcb
import coin.exchange.okex.kr_rest.native_private_client as nprivc
import coin.exchange.okex.kr_rest.product as exchange_product
from coin.exchange.base.kr_rest.rest_client_base import RestUpdate
from coin.exchange.okex.kr_rest.product import OkexCurrency

import coin.proto.coin_order_enums_pb2 as coin_order
from coin.proto.coin_query_pb2 import (AccountBalance,
                                       CurrencyBalance,
                                       ProductFillElement,
                                       ProductFillBundle,
                                       ProductOrderElement,
                                       ProductOrderBundle)


class OkexPrivateParser(object):
  ProductType = exchange_product.OkexProduct

  @staticmethod
  def parse_balance(json_msg):
    # Websocket balance update has no result field.
    assert json_msg.get('result') in (True, None), json_msg
    currency_balances = {}
    free_funds = None
    freezed_funds = None

    try:
      free_funds = json_msg['info']['funds']['free']
    except KeyError:
      pass
    try:  # For parsing websocket balance update message.
      free_funds = json_msg['info']['free']
    except KeyError:
      pass

    try:
      freezed_funds = json_msg['info']['funds']['freezed']
    except KeyError:
      pass
    try:  # For parsing websocket balance update message.
      freezed_funds = json_msg['info']['freezed']
    except KeyError:
      pass

    # Borrowed funds are ignored!
    for currency, amount in free_funds.items():
      # Sometimes Okex adds new coins, just ignore these unknown coins.
      currency = OkexCurrency.FromStrNativeCurrencyNothrow(currency)
      if currency is None:
        continue

      currency_balance = CurrencyBalance(currency=currency.currency,
                                         currency_native=currency.native_currency,
                                         available=float(amount))
      currency_balances[currency] = currency_balance

    if freezed_funds is not None:
      for currency, amount in freezed_funds.items():
        # Sometimes Okex adds new coins, just ignore these unknown coins.
        currency = OkexCurrency.FromStrNativeCurrencyNothrow(currency)
        if currency is None:
          continue

        if currency in currency_balances:
          currency_balances[currency].hold = float(amount)
        else:
          currency_balance = CurrencyBalance(currency=currency.currency,
                                             currency_native=currency.native_currency,
                                             hold=float(amount))
          currency_balances[currency] = currency_balance

    # Compute total balance.
    for balance in currency_balances.values():
      balance.total = balance.available + balance.hold

    return AccountBalance(exchange="Okex", each_balance=currency_balances.values())

  # status field meaning:
  # -1 = cancelled,
  # 0 = unfilled,
  # 1 = partially filled,
  # 2 = fully filled,
  # 3 = ???  (status 3 is not mentioned in the doc, but observed)
  # 4 = cancel request in process
  @staticmethod
  def parse_fill(update_msg, product) -> ProductFillBundle:
    """Only first 200 filled orders are retrieved."""
    assert update_msg['result']
    currency_fill_list = []
    for update_elem in update_msg['orders']:
      assert update_elem['status'] in (-1, 0, 1, 2, 3, 4), update_elem['status']
      if update_elem['status'] not in [1, 2]:
        continue
      if update_elem['symbol'] != product.native_symbol:
        continue

      if update_elem['type'] == 'sell':
        side = coin_order.SELL_ORDER
      elif update_elem['type'] == 'buy':
        side = coin_order.BUY_ORDER
      else:
        raise ValueError(update_elem['type'])

      currency_fill = ProductFillElement(symbol=product.symbol,
                                         side=side,
                                         price=float(update_elem['price']),
                                         qty=float(update_elem['deal_amount']),
                                         order_id=str(update_elem['order_id']))
      currency_fill_list.append(currency_fill)

    return ProductFillBundle(exchange='Okex', symbol=product.symbol, each_fill=currency_fill_list)

  @staticmethod
  def parse_orders(update_msg, product):
    """Only first 200 filled orders are retrieved."""
    assert update_msg['result']
    currency_order_list = []
    for update_elem in update_msg['orders']:
      assert update_elem['status'] in (-1, 0, 1, 2, 3, 4), update_elem['status']
      if update_elem['status'] not in [0, 1]:
        continue
      if update_elem['symbol'] != product.native_symbol:
        continue

      if update_elem['type'] == 'sell':
        side = coin_order.SELL_ORDER
      elif update_elem['type'] == 'buy':
        side = coin_order.BUY_ORDER
      else:
        raise ValueError(update_elem['type'])

      currency_order = ProductOrderElement(symbol=product.symbol,
                                           side=side,
                                           price=float(update_elem['price']),
                                           qty=float(update_elem['amount']),
                                           order_id=str(update_elem['order_id']),
                                           created_timestamp=int(
                                               float(update_elem['create_date']) * 1e+6))
      currency_order_list.append(currency_order)
    return ProductOrderBundle(exchange='Okex',
                              symbol=product.symbol,
                              each_order=currency_order_list)


class OkexPrivateClient(privcb.PrivateClientBase):
  ProductType = exchange_product.OkexProduct

  def __init__(self, key_file):
    privcb.PrivateClientBase.__init__(self)
    self.nprivc = nprivc.OkexNativePrivateClient(key_file=key_file)

  def query_account_balance_impl(self):
    update = self.nprivc.query_POST_userinfo()
    update.msg = OkexPrivateParser.parse_balance(update.msg)
    return update

  def query_fills_impl(self, product):
    page_length = 200  # 200 is the maximum page length.
    filled_status = 1

    native_msg = []
    fill_list = []
    current_page = 0
    symbol = product.native_symbol
    first_update = self.nprivc.query_POST_order_history(symbol,
                                                        status=filled_status,
                                                        current_page=current_page,
                                                        page_length=page_length)
    fill_list.extend(OkexPrivateParser.parse_fill(first_update.msg, product).each_fill)
    native_msg.append(first_update.msg)
    total_pages = math.ceil(first_update.msg['total'] / page_length)
    update = None
    for current_page in range(1, total_pages):
      update = self.nprivc.query_POST_order_history(symbol,
                                                    status=filled_status,
                                                    current_page=current_page,
                                                    page_length=page_length)
      fill_list.extend(OkexPrivateParser.parse_fill(update.msg, product).each_fill)
      native_msg.append(update.msg)

    currency_fill_bundle = ProductFillBundle(exchange='Okex',
                                             symbol=product.symbol,
                                             each_fill=fill_list)

    req_time = first_update.req_timestamp
    if update:
      res_time = update.res_timestamp
    else:
      res_time = first_update.res_timestamp
    merged_update = RestUpdate(req_time, res_time, currency_fill_bundle, native_msg=native_msg)
    return merged_update

  def query_submit_impl(self, product, side, price, qty):
    return self.nprivc.query_POST_trade(symbol=product.native_symbol,
                                        order_side=side,
                                        price=price,
                                        qty=qty)

  # order_id can be a list of ids or a single id.
  def query_cancel_impl(self, product, order_id):
    return self.nprivc.query_POST_cancel_order(symbol=product.native_symbol, order_id=order_id)

  def query_cancel_product_impl(self, product):
    orders = self.query_list_orders_impl(product)
    for an_order in orders.msg.each_order:
      self.query_cancel_impl(product, an_order.order_id)

  def query_list_orders_impl(self, product):
    symbol = product.native_symbol
    order_list = []
    first_update = self.nprivc.query_POST_order_info(symbol)
    order_list.extend(OkexPrivateParser.parse_orders(first_update.msg, product).each_order)
    req_time = first_update.req_timestamp
    res_time = first_update.res_timestamp
    currency_order_bundle = ProductOrderBundle(exchange='Okex',
                                               symbol=product.symbol,
                                               each_order=order_list)
    merged_update = RestUpdate(req_time,
                               res_time,
                               currency_order_bundle,
                               native_msg=first_update.msg)
    return merged_update

  def query_list_orders_impl_deprecated(self, product):
    page_length = 200  # 200 is the maximum page length.
    unfilled_status = 0

    native_msg = []
    order_list = []
    current_page = 0
    symbol = product.native_symbol
    first_update = self.nprivc.query_POST_order_history(symbol,
                                                        status=unfilled_status,
                                                        current_page=current_page,
                                                        page_length=page_length)
    order_list.extend(OkexPrivateParser.parse_orders(first_update.msg, product).each_order)
    native_msg.append(first_update.msg)
    total_pages = math.ceil(first_update.msg['total'] / page_length)

    update = None
    for current_page in range(1, total_pages):
      update = self.nprivc.query_POST_order_history(symbol,
                                                    status=unfilled_status,
                                                    current_page=current_page,
                                                    page_length=page_length)
      order_list.extend(OkexPrivateParser.parse_orders(update.msg, product).each_order)
      native_msg.append(update.msg)

    currency_order_bundle = ProductOrderBundle(exchange='Okex',
                                               symbol=product.symbol,
                                               each_order=order_list)

    req_time = first_update.req_timestamp
    if update:
      res_time = update.res_timestamp
    else:
      res_time = first_update.res_timestamp
    merged_update = RestUpdate(req_time, res_time, currency_order_bundle, native_msg=native_msg)
    return merged_update


class OkexExWalletPrivateClient(privcb.ExWalletPrivateClientBase):
  ProductType = exchange_product.OkexProduct

  def __init__(self, key_file):
    privcb.PrivateClientBase.__init__(self)
    self.nprivc = nprivc.OkexNativePrivateClient(key_file=key_file)

  def query_account_balance_impl(self):
    update = self.nprivc.query_POST_walletinfo()
    update.msg = OkexPrivateParser.parse_balance(update.msg)
    return update
